#!/usr/bin/env python3
"""
train_baseline.py
Train a simple RandomForest on prepared features. This assumes you created a labels.csv matching session start times.
"""
import argparse
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import joblib

parser = argparse.ArgumentParser()
parser.add_argument('--features', required=True)
parser.add_argument('--labels', required=True)
parser.add_argument('--out', default='models/rf_baseline.joblib')
args = parser.parse_args()

X = pd.read_parquet(args.features)
y = pd.read_csv(args.labels)
# assumes labels.csv has column 'start' matching features.start and 'label' (0/1)
merged = pd.merge(X, y, on='start')
label = merged['label']
features = merged.select_dtypes(include=[float, int]).drop(columns=['duration'])

X_train, X_test, y_train, y_test = train_test_split(features, label, test_size=0.3, random_state=42)
clf = RandomForestClassifier(n_estimators=200, class_weight='balanced')
clf.fit(X_train, y_train)

pred = clf.predict(X_test)
print(classification_report(y_test, pred))
joblib.dump(clf, args.out)
print('wrote model to', args.out)